# encoding: utf-8
# generated with chatGPT, but modified a lot by hand :-)
# not fully tested, used maybe need more idea from your own head

import pandas as pd
import gm.api as gm
import tushare as ts
import baostock as bs
import yfinance as yf

gm.set_serv_addr("MYQUANT_GM_SERVICE_IP:PORT")

gm.set_token("YOUR_TOKEN_MYQUANT")
pro = ts.pro_api('YOUR_TOKEN_TUSHARE')

def translate_symbol(symbol: str, data_source: str) -> str:
    exchange, symbol_code = symbol.split('.')
    if data_source == 'tushare':
        if exchange == 'SHSE':
            exchange_code = 'SH'
        elif exchange == 'SZSE':
            exchange_code = 'SZ'
        return f'{symbol_code}.{exchange_code}'
    elif data_source == 'baostock':
        if exchange == 'SHSE':
            exchange_code = 'sh'
        elif exchange == 'SZSE':
            exchange_code = 'sz'
        return f'{exchange_code}.{symbol_code}'
    elif data_source == 'yahoo':
        ticker_symbol = symbol.split('.')[1]
        ticker_symbol += '.SS' if symbol.startswith('SHSE') else '.SZ'
        return ticker_symbol

def get_close_prices(symbols: list, start_date: str, end_date: str, freq: str='d', data_source: str='myquant') -> pd.DataFrame:
    close_prices = pd.DataFrame()
    if data_source == 'myquant':
        for symbol in symbols:
            try:
                data = gm.history(symbol, frequency='1'+freq, start_time=start_date, end_time=end_date, fields='symbol,close,eob', df=True)
                if data.empty:
                    print(f'No data found for {symbol} using myquant.')
                else:
                    symbol_data = pd.DataFrame(data['close'].values, index=pd.to_datetime(data['eob']), columns=[symbol])
                    if close_prices.empty:
                        close_prices = symbol_data
                    else:
                        close_prices = close_prices.join(symbol_data)
            except Exception as e :
                 print(f"An error occurred while fetching historical prices from myquant. Error message :{e}")
    if data_source == 'tushare':
        for symbol in symbols:
            translated_symbol = translate_symbol(symbol, data_source)
            try:
                if freq=='d': 
                    data=pro.daily(ts_code=translated_symbol,start_date=start_date.replace('-', ''), end_date=end_date.replace('-', '')) 
                else:
                    print(f"Frequency {freq} is not supported by Tushare Pro.")
                    continue
                if data.empty:
                    print(f'No data found for {symbol} using Tushare Pro.')
                else:
                    symbol_data = pd.DataFrame(data['close'].values, index=pd.to_datetime(data['trade_date']), columns=[symbol])
                    symbol_data = symbol_data.iloc[::-1]
                    if close_prices.empty:
                        close_prices = symbol_data
                    else:
                        close_prices = close_prices.join(symbol_data)
            except Exception as e :
                 print(f"An error occurred while fetching historical prices from Tushare Pro. Error message :{e}")
    elif data_source == 'baostock':
        bs.login()
        for symbol in symbols:
            translated_symbol = translate_symbol(symbol, data_source)
            rs = bs.query_history_k_data_plus(translated_symbol,
                                              "date,close",
                                              start_date=start_date,
                                              end_date=end_date,
                                              frequency=freq)
            if rs.error_code != '0':
                print(f'Error retrieving data for {symbol} using Baostock: {rs.error_msg}')
            else:
                result = rs.get_data()
                if result.empty:
                    print(f'No data found for {symbol} using Baostock.')
                else:
                    symbol_data = pd.DataFrame(result['close'].values, index=pd.to_datetime(result['date']), columns=[symbol])
                    symbol_data[symbol] = pd.to_numeric(symbol_data[symbol])
                    if close_prices.empty:
                        close_prices = symbol_data
                    else:
                        close_prices = close_prices.join(symbol_data)
        bs.logout()
    elif data_source == 'yahoo':
        for symbol in symbols:
            translated_symbol = translate_symbol(symbol, data_source)
            try:
                ticker = yf.Ticker(translated_symbol)
                if freq=='d': 
                    interval='1d'
                elif freq=='w': 
                    interval='1wk'
                elif freq=='m': 
                    interval='1mo'           
                
                data = ticker.history(start=start_date, end=end_date,interval=interval)
                if not len(symbol_data):
                    print(f'No historical price found for {translated_symbol}.')
                else:
                    symbol_data = pd.DataFrame(data['close'].values, index=pd.to_datetime(data['trade_date']), columns=[symbol])
                    if close_prices.empty:
                        close_prices = symbol_data
                    else:
                        close_prices = close_prices.join(symbol_data)
                
            except Exception as e :
                 print(f"An error occurred while fetching historical prices from Yahoo Finance. Error message :{e}")

    return close_prices